import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from python_ai.common.xcommon import *

np.random.seed(3)
HOW_MANY_TO_TEST = 200

vcode_dir = '../../../../large_data/DL1/_many_files/vcode_data/test'
names = os.listdir(vcode_dir)
n_names = len(names)
names = np.array(names)
idx = np.random.permutation(n_names)  # for shuffle
names = names[idx][:HOW_MANY_TO_TEST]
x = []
y = []
for name in names:
    digits, ext = os.path.splitext(name)
    if ext.lower() != '.jpg':
        continue
    digits_list = list(digits)
    if len(digits_list) != 4:
        continue
    digits_arr = []
    try:
        for di in digits_list:
            digits_arr.append(int(di))
    except ValueError as ex:
        continue

    y.append(digits_arr)

    path = os.path.join(vcode_dir, name)
    img = cv.imread(path, cv.IMREAD_GRAYSCALE)
    x.append(img)

x = np.uint8(x)
y = np.uint8(y)
pred = np.zeros_like(y, dtype=np.uint8)
LEN_X = len(x)

print_numpy_ndarray_info(x, 'x')
print_numpy_ndarray_info(y, 'y')

spr = 6  # 6
spc = 10  # 10
spn = 0
plt.figure(figsize=[14, 7])

def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    if spn > spr * spc:
        return
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


def digits_arr_to_str(digits_arr):
    return ''.join([str(di) for di in digits_arr.tolist()])


import python_ai.CV_2.project.vcode.tf2x_lenet_mnist_by_input_output as saved

saved.model.load_weights(r'_save\tf2x_lenet_mnist_by_input_output.py\2.1_0.001_8_128\weights')

kernel = np.ones((3, 3), dtype=np.uint8)
SIDE = 50
ptsb_1 = np.float32(((0, 0), (SIDE - 1, 0), (SIDE - 1, SIDE - 1)))
ptsb_4_block = np.float32(((0, 0), (SIDE * 4 - 1, 0), (SIDE * 4 - 1, SIDE - 1)))
ptsb_4_split_arr = (
     np.float32(((0, 0), ((SIDE - 1), 0), ((SIDE - 1), SIDE - 1))),
     np.float32((((SIDE - 1), 0), ((SIDE - 1) * 2, 0), ((SIDE - 1) * 2, SIDE - 1))),
     np.float32((((SIDE - 1) * 2, 0), ((SIDE - 1) * 3, 0), ((SIDE - 1) * 3, SIDE - 1))),
     np.float32((((SIDE - 1) * 3, 0), ((SIDE - 1) * 4, 0), ((SIDE - 1) * 4, SIDE - 1))),
)
ptsb_3_block = np.float32(((0, 0), (SIDE * 3 - 1, 0), (SIDE * 3 - 1, SIDE - 1)))
ptsb_3_split_arr = (
     np.float32(((0, 0), ((SIDE - 1), 0), ((SIDE - 1), SIDE - 1))),
     np.float32((((SIDE - 1), 0), ((SIDE - 1) * 2, 0), ((SIDE - 1) * 2, SIDE - 1))),
     np.float32((((SIDE - 1) * 2, 0), ((SIDE - 1) * 3, 0), ((SIDE - 1) * 3, SIDE - 1))),
)
ptsb_2_block = np.float32(((0, 0), (SIDE * 2 - 1, 0), (SIDE * 2 - 1, SIDE - 1)))
ptsb_2_split_arr = (
     np.float32(((0, 0), ((SIDE - 1), 0), ((SIDE - 1), SIDE - 1))),
     np.float32((((SIDE - 1), 0), ((SIDE - 1) * 2, 0), ((SIDE - 1) * 2, SIDE - 1))),
)
SIDE2 = 28
PADDING = 4
pts_semi = np.float32(((0, 0), (SIDE2 - 1, 0), (SIDE2 - 1, SIDE2 - 1)))
pts_final = np.float32(((PADDING, PADDING), (SIDE2 - 1 - PADDING, PADDING), (SIDE2 - 1 - PADDING, SIDE2 - 1 - PADDING)))
good_pred_idx = []
idx_cate_by_n_block = [[], [], [], []]
idx4iter = np.arange(LEN_X)
sample_idx = [14, 60, 13, 19, 46, 42, 27, 5]
idx4iter[:len(sample_idx)] = sample_idx
for i in idx4iter:
    sep(f'#{i + 1}')
    # show
    label = digits_arr_to_str(y[i])
    my_show_img(x[i], f'#{i + 1}: ' + label, cmap='gray')

    img = x[i].copy()

    # erode
    ret, img = cv.threshold(img, 0, 255, cv.THRESH_OTSU + cv.THRESH_BINARY_INV)
    # img = cv.erode(img, kernel, iterations=1)
    img = cv.morphologyEx(img, cv.MORPH_OPEN, kernel, iterations=1)
    # my_show_img(img, 'open', cmap='gray')

    # contours
    contours, _ = cv.findContours(img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    n_contours = len(contours)

    area_arr = []
    pts_arr = []
    pts_full_arr = []
    left_side_arr = []
    width_arr = []
    for con in contours:
        rotated = cv.minAreaRect(con)
        # the rotation angle. When the angle is 0, 90, 180, 270 etc., the rectangle becomes an up-right rectangle.
        (cx, cy), (w, h), angle = rotated
        left_side_arr.append(cx)
        pts = cv.boxPoints(rotated)
        area = cv.contourArea(pts)
        area_arr.append(area)
        pts_full_arr.append(pts)
        if angle > 45:
            pts_arr.append(pts[:3])
            width_arr.append(w)
        else:
            pts_arr.append(pts[1:])
            width_arr.append(h)
    # build numpy array
    area_arr = np.int32(area_arr)
    pts_arr = np.float32(pts_arr)
    pts_full_arr = np.float32(pts_full_arr)
    left_side_arr = np.int32(left_side_arr)
    width_arr = np.int32(width_arr)
    # sort by area
    idx_top4 = area_arr.argsort()
    idx_top4 = idx_top4[::-1][:4]  # Top 4 area
    # select pts
    pts_arr = pts_arr[idx_top4]
    pts_full_arr = pts_full_arr[idx_top4]
    # select left side
    left_side_arr = left_side_arr[idx_top4]
    width_arr = width_arr[idx_top4]
    # select area
    area_arr = area_arr[idx_top4]

    # filter by area
    area_mean = area_arr.mean()
    n_blocks = 0
    factor = 5
    while not 1 <= n_blocks <= 4:
        idx_selected = []
        for ii, area in enumerate(area_arr):
            if area > area_mean / factor:
                idx_selected.append(ii)
        n_blocks = len(idx_selected)
        factor += 0.5

    # select pts
    pts_arr = pts_arr[idx_selected]
    pts_full_arr = pts_full_arr[idx_selected]
    # select left side
    left_side_arr = left_side_arr[idx_selected]
    width_arr = width_arr[idx_selected]
    width_mean = width_arr.mean()
    width_max_i = width_arr.argmax()
    # select area
    area_arr = area_arr[idx_selected]
    area_mean = area_arr.mean()
    area_max_i = area_arr.argmax()

    # get order
    idx_ordered = left_side_arr.argsort()
    # apply order
    pts_arr = pts_arr[idx_ordered]
    pts_full_arr = pts_full_arr[idx_ordered]

    # digits
    img_color = cv.cvtColor(img, cv.COLOR_GRAY2BGR)
    x_data = []
    print(n_blocks)
    idx_cate_by_n_block[n_blocks - 1].append(i)
    if n_blocks == 2:
        is_half_half = np.isclose(width_arr[0], width_arr[1], atol=width_mean * 0.2)
    for ii, pts in enumerate(pts_arr):
        width = width_arr[ii]
        area = area_arr[ii]

        def process(pts2a=None):
            if pts2a is not None:
                M = cv.getAffineTransform(pts2a, pts_semi)
                img_di = cv.warpAffine(img_block, M, (SIDE2, SIDE2))
            else:
                img_di = img_block.copy()
            M = cv.getAffineTransform(pts_semi, pts_final)
            img_di = cv.warpAffine(img_di, M, (SIDE2, SIDE2))
            my_show_img(img_di, 'digit')
            x_data.append(img_di)
            pts4poly = np.expand_dims(pts_full_arr[ii], axis=(0, 1)).astype(np.int32)
            cv.polylines(img_color, pts4poly, True, (0, 255, 0), 2)

        if n_blocks == 1:
            M = cv.getAffineTransform(pts, ptsb_4_block)
            img_block = cv.warpAffine(img, M, (SIDE * 4, SIDE))
            for pts2a in ptsb_4_split_arr:
                process(pts2a)
        elif n_blocks == 2:
            if is_half_half:
                M = cv.getAffineTransform(pts, ptsb_2_block)
                img_block = cv.warpAffine(img, M, (SIDE * 2, SIDE))
                for pts2a in ptsb_2_split_arr:
                    process(pts2a)
            else:
                if width > width_mean:
                    M = cv.getAffineTransform(pts, ptsb_3_block)
                    img_block = cv.warpAffine(img, M, (SIDE * 3, SIDE))
                    for pts2a in ptsb_3_split_arr:
                        process(pts2a)
                else:
                    M = cv.getAffineTransform(pts, ptsb_1)
                    img_block = cv.warpAffine(img, M, (SIDE, SIDE))
                    M = cv.getAffineTransform(ptsb_1, pts_semi)
                    img_block = cv.warpAffine(img_block, M, (SIDE2, SIDE2))
                    process()
        elif n_blocks == 3:
            if width_max_i == ii:
                M = cv.getAffineTransform(pts, ptsb_2_block)
                img_block = cv.warpAffine(img, M, (SIDE * 2, SIDE))
                for pts2a in ptsb_2_split_arr:
                    process(pts2a)
            else:
                M = cv.getAffineTransform(pts, ptsb_1)
                img_block = cv.warpAffine(img, M, (SIDE, SIDE))
                M = cv.getAffineTransform(ptsb_1, pts_semi)
                img_block = cv.warpAffine(img_block, M, (SIDE2, SIDE2))
                process()
        else:
            M = cv.getAffineTransform(pts, ptsb_1)
            img_block = cv.warpAffine(img, M, (SIDE, SIDE))
            M = cv.getAffineTransform(ptsb_1, pts_semi)
            img_block = cv.warpAffine(img_block, M, (SIDE2, SIDE2))
            process()

    x_data = np.uint8(x_data).astype(np.float32) / 255.
    predict = saved.model.predict(x_data).argmax(axis=1)
    title = digits_arr_to_str(predict)
    my_show_img(img_color, label + ': ' + title, lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
    predict = predict[:4]
    pred[i] = predict
    if title == label:
        good_pred_idx.append(i)

acc_digits = (pred == y).mean()
print(f'Accuracy of digits: {acc_digits}')
len_good = len(good_pred_idx)
print(f'Good predictions: {len_good}, accuracy = {len_good / LEN_X} (theory: {acc_digits**4})')
print(good_pred_idx)
print('Category by n blocks:')
print(idx_cate_by_n_block)

plt.show()
